from flask import Flask, request
import requests
import time
import os
import jwt


app = Flask(__name__)

API_KEY: str = os.getenv("API_KEY", "c0778e3af6877d5081b7c03ef0e35bf7.P3NPp6ikDA3T2U1I")
# 设置请求的URL
url = 'https://open.bigmodel.cn/api/paas/v3/model-api/charglm-3/invoke'


@app.route('/api/sse', methods=['POST'])
def sse():
    req = request.json
    # 设置请求参数
    params = req

    # 发起流式POST请求
    response = requests.post(
            url,
            headers= {
                'Content-Type': 'application/json',
                'Authorization': generate_token(API_KEY, 1800)
            },
            json=params)
    return response.content.decode('utf-8')


def generate_token(apikey: str, exp_seconds: int) -> str:
    print('apikey:', apikey)
    try:
        id, secret = apikey.split(".")
    except Exception as e:
        raise Exception("invalid apikey", e)

    payload = {
        "api_key": id,
        "exp": int(round(time.time() * 1000)) + exp_seconds * 1000,
        "timestamp": int(round(time.time() * 1000)),
    }

    return jwt.encode(
        payload,
        secret,
        algorithm="HS256",
        headers={"alg": "HS256", "sign_type": "SIGN"},
    )

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8006)